package xyz.syyrjx.blog.util.ai;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.aigc.generation.models.QwenParam;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.ResultCallback;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.NoApiKeyException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import xyz.syyrjx.blog.entity.ai.AiDialogue;
import xyz.syyrjx.blog.exception.AiException;
import xyz.syyrjx.blog.util.CachePool;
import xyz.syyrjx.blog.util.ConstValuePool;
import xyz.syyrjx.blog.util.StringUtil;
import xyz.syyrjx.blog.websocket.AiWebsocketService;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.Semaphore;

/**
 * @author 大地崩坏苍蝇兽
 * @date 2024/2/18 12:48
 * @description 通义千问模型服务
 * @since v0.0.1
 */
@Component
public class QwenModelService implements AiModelService{

    private Generation gen;

    @Resource
    private AiWebsocketService aiWebsocketService;

    public void createGen(){
        gen = new Generation();
    };

    private static final Logger logger = LoggerFactory.getLogger(QwenModelService.class);
    @Override
    public String answer(String prompt, HttpServletRequest request, String identity) {

        List<AiDialogue> dialogues = CachePool.AI_DIALOGUE_LIST_MAP.get(identity)
                .computeIfAbsent(ConstValuePool.QWEN_DIALOGUES, k -> new LinkedList<>());
        dialogues.add(AiDialogue.createUserDialogue(prompt));

        List<Message> msgManager = CachePool.QWEN_MESSAGE_DIALOGUES_MAP.get(identity);
        QwenParam param = CachePool.QWEN_PARAM_MAP.get(identity);
        if (msgManager == null) {
            msgManager = new ArrayList<>();
            CachePool.QWEN_MESSAGE_DIALOGUES_MAP.put(identity, msgManager);
            Message systemMsg = Message.builder()
                    .role(Role.SYSTEM.getValue())
                    .content("You are a helpful assistant.")
                    .build();
            msgManager.add(systemMsg);
            Message userMsg = Message.builder()
                    .role(Role.USER.getValue())
                    .content(prompt)
                    .build();
            msgManager.add(userMsg);
        }else {
            msgManager.add(Message.builder().role("user").content(prompt).build());
            param.setMessages(msgManager);
        }
        if (param == null) {
            param = QwenParam.builder()
                    .model(Generation.Models.QWEN_MAX)
                    .messages(msgManager)
                    .resultFormat(QwenParam.ResultFormat.MESSAGE)
                    .topP(0.8)
                    .enableSearch(true)
                    .incrementalOutput(true)
                    .build();
            CachePool.QWEN_PARAM_MAP.put(identity, param);
        }


        try {
            logger.debug("发送的请求为{}",param);
            Semaphore semaphore = new Semaphore(0);
            StringBuilder resultBuilder = new StringBuilder();
            gen.streamCall(param, new ResultCallback<GenerationResult>(){

                @Override
                public void onEvent(GenerationResult generationResult) {
                    String newMessage = generationResult.getOutput().getChoices().get(0).getMessage().getContent();
                    StringBuilder finalResBuilder = resultBuilder.append(newMessage);
                    if (1 == (1 & StringUtil.countSubStr(finalResBuilder,ConstValuePool.MARKDOWN_CODE_BLOCK_START))) {
                        finalResBuilder = new StringBuilder(finalResBuilder)
                                .append(ConstValuePool.MARKDOWN_CODE_BLOCK_END);
                    }
                    aiWebsocketService.sendMessage(finalResBuilder.toString(), identity);
                }

                @Override
                public void onComplete() {
                    semaphore.release();
                }

                @Override
                public void onError(Exception e) {
                    semaphore.release();
                    logger.error("通义千问运行出错, 报错栈如下");
                    Throwable t = e;
                    while (t != null) {
                        logger.error( t.toString());
                        t = e.getCause();
                    }
                }
            });
            semaphore.acquire();
            String resString = resultBuilder.toString();
            msgManager.add(Message.builder().role("assistant").content(resString).build());
            while (msgManager.size() > ConstValuePool.QWEN_MAX_MESSAGE
                    || !"user".equals(msgManager.get(1).getRole())) {
                msgManager.remove(1);
            }
            AiLogUtil.logDialogues(logger,"通义千问",prompt,resString,request);
            dialogues.add(AiDialogue.createAssistantDialogue(resString));
            return "";
        } catch (NoApiKeyException e) {
            logger.error("调用通义千问缺少ApiKey");
            throw new AiException("没有ApiKey", e);
        } catch (Exception e) {
            logger.error("调用通义千问出现问题：{}",e.getMessage());
            throw new AiException("出现了一些问题", e);
        }
    }

}
